import pandas
import sys
import os,re
sys.path.append(r"D:\ProgramData\anaconda3\envs\sparktts\Spark-TTS")
import torch
import soundfile as sf
import logging
from datetime import datetime
import platform
from cli.SparkTTS import SparkTTS
from pydub import AudioSegment
from pydub.silence import detect_silence

##############################################
# pitch  and speed range
# very_low | low | moderate | high | very_high
##############################################

class spark_batch_generate():

    vioce_param_map = {
        1.0: "very_low",
        2.0: "low",
        3.0: "moderate",
        4.0: "high",
        5.0: "very_high"
    }
    def __init__(self):
        self._init_tts_model()

    @classmethod
    def shorten_silence(self, input_file):
        """
        去除音频文件中的长静音片段。
        """
        # 读取音频文件
        audio = AudioSegment.from_file(input_file)
        # 检测静音片段，阈值为 -40 dBFS，最小静音时长为 100 毫秒
        silence_ranges = detect_silence(audio, min_silence_len=100, silence_thresh=-40)
        new_audio = AudioSegment.empty()
        start_point = 0
        for start, end in silence_ranges:
            # 如果静音时长超过 3 秒
            if (end - start) > 3000:
                new_audio += audio[start_point:start]
                # 只保留 1 秒的静音
                new_audio += AudioSegment.silent(duration=1000)
                start_point = end
                print(">>>>>>Silence detected ： %s"%input_file)
            else:
                new_audio += audio[start_point:end]
                start_point = end
        new_audio += audio[start_point:]
        # 导出处理后的音频文件
        return new_audio

    def read_voice_data(self,seed_path,sheet_name):
        #读取音色数据excel文件,生成音色数据字典
        voice_data = pandas.read_excel(seed_path, sheet_name=sheet_name, header=0)
        keys=voice_data.columns[0] #获取第一列
        voice_data = voice_data.set_index(keys)
        voice_dict = voice_data.to_dict('index') #将数据转换为字典
        print('音色数据读取完成。')
        self.voice_dict = voice_dict
        return voice_dict

    def read_text_data(self,xls_path,sheet_name):
        #读取文本数据excel文件,生成文本数据字典
        text_data = pandas.read_excel(xls_path, sheet_name=sheet_name, header=0)
        text_dict = text_data.to_dict(orient='records') #将数据转换为字典
        print('文本数据读取完成。')
        return text_dict

    def _init_tts_model(self):
        #模型路径
        model_dir = r"D:\ProgramData\anaconda3\envs\sparktts\Spark-TTS\pretrained_models\Spark-TTS-0.5B"
        #gpu设备
        device = "0"
        logging.info(f"Using model from: {model_dir}")
        # Convert device argument to torch.device
        if platform.system() == "Darwin" and torch.backends.mps.is_available():
            # macOS with MPS support (Apple Silicon)
            device = torch.device(f"mps:{device}")
            logging.info(f"Using MPS device: {device}")
        elif torch.cuda.is_available():
            # System with CUDA support
            device = torch.device(f"cuda:{device}")
            logging.info(f"Using CUDA device: {device}")
        else:
            # Fall back to CPU
            device = torch.device("cpu")
            logging.info("GPU acceleration not available, using CPU")

        # Initialize the model
        self.model = SparkTTS(model_dir, device)


    def run_tts(self,vocice_info,text_info,save_dir,index):
        
        # Generate unique filename using timestamp
        # timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        save_path = os.path.join(save_dir, f"{index+2}.wav")

        logging.info("Starting inference...")

        # Perform inference and save the output audio
        # print(text_info["pitch"])
        try:
            pitch = self.vioce_param_map[text_info["pitch"]]
        except KeyError:
            pitch =vocice_info["pitch"]
        try:
            speed = self.vioce_param_map[text_info["speed"]]
        except KeyError:
            speed =vocice_info["speed"]
        
        # 重新获取prompt_speech_path和prompt_text
        volume = text_info["volume"]
        actor_wavs = vocice_info["prompt_speech_path"]
        wavs_list = os.listdir(actor_wavs)
        wavs_list = [os.path.join(actor_wavs, wav) for wav in wavs_list if wav.endswith(".wav") or wav.endswith(".mp3")]
        if isinstance(volume,str):
            for wav in wavs_list:
                if re.search(volume,wav):
                    prompt_speech_path = wav
                    index_seed = re.search(r"seed_(\d+)",wav).group(1)
                    prompt_text = self.voice_dict[int(index_seed)]["prompt_text"]
                    break
                elif re.search("default",wav):
                    default_wav = wav
            else:
                print("没有找到对应的音频文件: %s"%volume)
                prompt_speech_path = default_wav
                prompt_text = vocice_info["prompt_text"]
        else:
            for wav in wavs_list:
                if re.search("default",wav):
                    prompt_speech_path = wav
                    break
            prompt_text = vocice_info["prompt_text"]


        print(text_info["text"]) 
        print(prompt_speech_path)
        print(prompt_text)
        print(pitch)
        print(speed)

        #推理    
        with torch.no_grad():

            wav = self.model.inference(
                text = text_info["text"],
                prompt_speech_path=prompt_speech_path,
                prompt_text=prompt_text,
                pitch=pitch,
                speed=speed,
                temperature=0.95,
                top_k=30,
                top_p=0.95,
            )

            sf.write(save_path, wav, samplerate=16000)
        print(f"Audio saved at: {save_path}")
        logging.info(f"Audio saved at: {save_path}")

    def batch_run(self,voice_dict,text_dict,save_dir,index=0):
        #需要从失败处重新开始的话，可以设置index，index=excel行号-2
        # Ensure the save directory exists
        os.makedirs(save_dir, exist_ok=True)
        #批量生成音频
        for i,text_info in enumerate(text_dict):
            if i >= index:                              #从指定行开始
                voice_info  = voice_dict[text_info["seed"]]
                self.run_tts(voice_info,text_info,save_dir,i)

    @classmethod
    def splice_wav(cls,directory, output_file):
        # 初始化一个空的 AudioSegment 对象
        combined = AudioSegment.empty()
        # 遍历目录中的所有文件

        for filename in sorted(os.listdir(directory),key=cls.natural_keys):
            if filename.endswith('.wav'):
                file_path = os.path.join(directory, filename)
                try:
                    # 加载音频文件
                    # audio = AudioSegment.from_wav(file_path)
                    audio = cls.shorten_silence(file_path)
                    # 拼接音频
                    combined += audio
                except Exception as e:
                    print(f"Error processing {filename}: {e}")
        # 保存拼接后的音频文件
        if len(combined) > 0:
            combined.export(output_file, format='wav')
            print(f"Audio files concatenated and saved to {output_file}")

    @classmethod
    def atoi(cls,text):
        return int(text) if text.isdigit() else text

    @classmethod
    def natural_keys(cls,text):
        return [cls.atoi(c) for c in re.split(r'(\d+)', text)]

if __name__ == "__main__":
    spark_batch_generate = spark_batch_generate()
    seed_path = r'E:\txt\seed2.xlsx'
    sheet_name = 'sparktts'
    voice_dict = spark_batch_generate.read_voice_data(seed_path,sheet_name)

    file_name = "第一章_翘臀小美女"
    sheet_name = 'Sheet'
    xls_path = r'E:\txt\通天之路\通天之路\%s.xls'%file_name
    text_dict = spark_batch_generate.read_text_data(xls_path,sheet_name)

    # #存储路径
    save_dir = r"E:\txt\wav"
    spark_batch_generate.batch_run(voice_dict,text_dict,save_dir,index=76) 
    spark_batch_generate.splice_wav(r"E:\txt\wav", r"E:\txt\%s.wav"%file_name)

    # voice_dict={
    #     1: {
    #         "prompt_speech_path": r"E:\txt\音频资源\高叶.MP3",
    #         "prompt_text": "在京海，掉下个钢镚儿都得姓高。",
    #         "pitch": "very_low",
    #         "speed": "moderate"
    #     },
    #     2: {
    #         "prompt_speech_path": r"E:\txt\音频资源\seed\seed_2.wav",
    #         "prompt_text": "这是一段测试文本,我是杨幂。",
    #         "pitch": "moderate",
    #         "speed": "moderate"
    #     }
    # }
    # text_info = {"text":"这是一段测试文本,我是杨幂。祝大家天天开心。","pitch":1.0,"speed":1.0}
    # spark_batch_generate.run_tts(vocice_info=voice_dict[22],text_info=text_info,save_dir=r"E:\txt\wav",index=4)
